# 多波束合成: 自编码器实现自监督学习
import argparse
import logging
import numpy as np

from util.util_log import setup_logging
from util.util_image import draw_img_xyz_2, draw_img

import torch
import torch.nn as nn
import torch.optim as optim

# ================================================= 天线计算 Tensor =================================================
# 添加一个小常数以避免 log(0)
eps = 1e-8

# 计算距离
def distance(p1, p2):
    return torch.sqrt(torch.sum((p1 - p2) ** 2))

# 计算方向图, fft方式, xyz坐标
def get_pattern_fft(phase, posx, posy, k, Feed):
    Nx, Ny = len(posx), len(posy)
    Kfft, Lfft = 360, 361
    # 生成相位矩阵，这里假设已经得到了正确的相位值
    phase_matrix = torch.exp(1j * phase)
    # 根据波长和阵元位置计算u和v
    u = torch.linspace(-1, 1, Kfft)
    v = torch.linspace(-1, 1, Lfft)
    u, v = torch.meshgrid(u, v, indexing='ij')
    # 计算每个单元到馈源的距离差分项
    distance_diff = torch.zeros(Nx, Ny)
    for i in range(Nx):
        for j in range(Ny):
            E = torch.tensor([posx[i], posy[j], 0.0])
            rfij = distance(Feed, E)
            distance_diff[i, j] = -k * rfij
    # 将距离差分项加到相位矩阵中
    phase_matrix *= torch.exp(1j * distance_diff)
    # 使用FFT进行快速计算
    IF = torch.fft.ifftshift(torch.fft.ifft2(phase_matrix, s=(Kfft, Lfft)))
    F = IF * Kfft * Lfft  # 这里乘以采样点数是为了匹配直接计算的结果尺度
    # 返回绝对值作为方向图
    return torch.abs(F), u, v

# 相位生成方向图, xy坐标系, FFT方式
def phase_2_pattern_xyz_fft_torch(phaseBit):
    # 初始化参数
    nRow = 64  # x方向单元个数，θ方向
    mLine = 64  # y方向单元个数，φ方向
    f0 = 300  # 工作频点，单位：GHz
    c = 3e8  # 光速，单位：m/s
    lamda0 = c / (f0 * 1e9)  # 波长，单位：m
    dx = lamda0 / 2  # x方向单元间距，单位：m
    dy = dx  # y方向单元间距，单位：m
    k = 2 * np.pi / lamda0  # 波数
    # 生成天线位置坐标
    posx = torch.arange(-dx * (nRow - 1) / 2, dx * (nRow - 1) / 2 + dx, dx)
    posy = torch.arange(-dy * (mLine - 1) / 2, dy * (mLine - 1) / 2 + dy, dy)
    # 计算馈元位置
    Feed = torch.tensor([0.0, 0, 0.03862741699796952])
    # 计算方向图
    pattern, th, ph = get_pattern_fft(phaseBit, posx, posy, k, Feed)
    # 转换为直角坐标系
    x = th
    y = ph
    return pattern, x, y

def get_patterndbw_by_phase_torch(phaseRad):
    # 计算方向图
    patternBit, x, y = phase_2_pattern_xyz_fft_torch(phaseRad)  # FFT法计算, 快速
    # pattern 转dB
    pattern_dbw = 20 * torch.log10(torch.abs(patternBit) / torch.max(torch.max(torch.abs(patternBit))) + eps)
    return pattern_dbw

def loss_function_mask(output_add, AM1, AM2):
    alpha, beta = 1000, 1
    # 计算方向图
    A_add = get_patterndbw_by_phase_torch(output_add)
    # 计算峰值保留损失
    peak_loss = torch.sum(torch.abs(A_add * AM1) ** 2) + torch.sum(torch.abs(A_add * AM2) ** 2)
    # 计算背景抑制损失
    background_loss = torch.sum(torch.abs(A_add * (1 - AM1)) ** 2) + torch.sum(torch.abs(A_add * (1 - AM2)) ** 2)
    # 计算总损失
    total_loss = alpha * peak_loss + beta * background_loss
    print(f"total_loss={total_loss.item()}, peak_loss={peak_loss.item()}, background_loss={background_loss.item()}")
    return total_loss

# ================================================= 天线计算 numpy ==================================================
# from math import sin, cos, radians, sqrt
#
# # 计算距离
# def distance(p1, p2):
#     return sqrt(sum((np.array(p1) - np.array(p2)) ** 2))
#
# # 计算方向图, fft方式, xyz坐标
# def get_pattern_fft(phase, posx, posy, k, Feed):
#     Nx, Ny = len(posx), len(posy)
#     Kfft, Lfft = 360, 361
#     # Kfft, Lfft = 720, 720
#     # 生成相位矩阵，这里假设已经得到了正确的相位值
#     phase_matrix = np.exp(1j * phase)
#     # 根据波长和阵元位置计算u和v
#     u = np.linspace(-1, 1, Kfft)
#     v = np.linspace(-1, 1, Lfft)
#     u, v = np.meshgrid(u, v, indexing='ij')
#     # 计算每个单元到馈源的距离差分项
#     distance_diff = np.zeros((Nx, Ny))
#     for i in range(Nx):
#         for j in range(Ny):
#             E = np.array([posx[i], posy[j], 0])
#             rfij = distance(Feed, E)
#             distance_diff[i, j] = -k * rfij
#     # 将距离差分项加到相位矩阵中
#     phase_matrix *= np.exp(1j * distance_diff)
#     # 使用FFT进行快速计算
#     IF = np.fft.fftshift(np.fft.ifft2(phase_matrix, s=(Kfft, Lfft)))
#     F = IF * Kfft * Lfft  # 这里乘以采样点数是为了匹配直接计算的结果尺度
#     # 返回绝对值作为方向图
#     return np.abs(F), u, v  # 注意这里的u, v代替了原来的th, ph
#
# # 相位生成方向图, xy坐标系, FFT方式
# def phase_2_pattern_xyz_fft(phaseBit):
#     # 初始化参数
#     nRow = 64  # x方向单元个数，θ方向
#     mLine = 64  # y方向单元个数，φ方向
#     f0 = 300  # 工作频点，单位：GHz
#     c = 3e8  # 光速，单位：m/s
#     lamda0 = c / (f0 * 1e9)  # 波长，单位：m
#     dx = lamda0 / 2  # x方向单元间距，单位：m
#     dy = dx  # y方向单元间距，单位：m
#     k = 2 * np.pi / lamda0  # 波数
#     # 1. 相位生成方向图, 极坐标
#     # 生成天线位置坐标
#     posx = np.arange(-dx * (nRow - 1) / 2, dx * (nRow - 1) / 2 + dx, dx)
#     posy = np.arange(-dy * (mLine - 1) / 2, dy * (mLine - 1) / 2 + dy, dy)
#     # 计算馈元位置
#     Feed = [0.0, 0, 0.03862741699796952]
#     # 计算方向图
#     pattern, th, ph = get_pattern_fft(phaseBit, posx, posy, k, Feed)
#     # 2. 转换为直角坐标系
#     x = th
#     y = ph
#     return pattern, x, y
#
# def get_patterndbw_by_phase(phaseRad):
#     # 计算方向图
#     patternBit, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
#     # pattern 转dB
#     pattern_dbw = 20 * np.log10(np.abs(patternBit) / np.max(np.max(np.abs(patternBit))) + eps)
#     return pattern_dbw
#
# def loss(output_add, AM1, AM2):
#     alpha, beta = 1000, 1
#     # 计算方向图
#     A_add = get_patterndbw_by_phase(output_add)
#     # 计算峰值保留损失
#     peak_loss = np.sum(np.abs(A_add * AM1) ** 2) + np.sum(np.abs(A_add * AM2) ** 2)
#     # 计算背景抑制损失
#     background_loss = np.sum(np.abs(A_add * (1 - AM1)) ** 2) + np.sum(np.abs(A_add * (1 - AM2)) ** 2)
#     # 计算总损失
#     total_loss = alpha * peak_loss + beta * background_loss
#     print(f"total_loss={total_loss}, peak_loss={peak_loss}, background_loss={background_loss}")
#     return total_loss

# ================================================= 天线计算相关 ======================================================
from multi_beam.multi_beam_PS_complex import psm_complex_beam_n
from collections import deque
from util.util_ris_pattern import phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft, \
    phase_2_pattern, point_2_phase, eps


def get_patterndbw(phaseRad, bit_num=1):
    # 计算方向图
    # patternBit = phase_2_pattern(phaseRad)                # 公式法直接计算, 准确但速度太慢
    patternBit, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
    # pattern 转dB
    pattern_dbw = 20 * np.log10(np.abs(patternBit) / np.max(np.max(np.abs(patternBit))) + eps)
    return pattern_dbw


def find_peak_and_mark(patternBit_dbw_mask):
    # 获取数组的形状
    rows, cols = patternBit_dbw_mask.shape
    # 找到最大值及其索引
    # max_value = np.max(patternBit_dbw_mask)
    max_index = np.unravel_index(np.argmax(patternBit_dbw_mask), patternBit_dbw_mask.shape)
    # 初始化结果数组
    result = np.zeros_like(patternBit_dbw_mask, dtype=int)
    result[max_index] = 1
    # 使用队列进行广度优先搜索
    queue = deque([max_index])
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    while queue:
        current_row, current_col = queue.popleft()

        for dr, dc in directions:
            new_row, new_col = current_row + dr, current_col + dc

            if 0 <= new_row < rows and 0 <= new_col < cols:
                if patternBit_dbw_mask[new_row, new_col] < patternBit_dbw_mask[current_row, current_col]:
                    result[new_row, new_col] = 1
                    queue.append((new_row, new_col))

    return result


def get_mask_by_slope(bit_num, phaseRad):
    patternBit_dbw = get_patterndbw(phaseRad, bit_num)
    patternBit_dbw_mask = np.copy(patternBit_dbw)
    # print("ggggggggggggggggggggggggggggggggggggggggggggggggggggggggg")
    # draw_img(patternBit_dbw_mask)
    patternBit_dbw_mask = find_peak_and_mark(patternBit_dbw_mask)
    return patternBit_dbw, patternBit_dbw_mask

def get_traces(theta_start, theta_step, theta_end, phi_start, phi_step, phi_end):
    """ 获取波束指向和对应码阵 """
    logger.info(f"theta_start={theta_start}, theta_step={theta_step}, theta_end={theta_end}, "
                f"phi_start={phi_start}, phi_step={phi_step}, phi_end={phi_end}")
    traces = []
    for phi in range(phi_start, phi_end, phi_step):
        # 生成轨迹指向
        traces_idx = []
        for theta in range(theta_start, theta_end + 1, theta_step):
            traces_idx.append([theta, phi])
        # 记录
        traces.extend(traces_idx)
    logger.info(f"len of traces: {len(traces)}")
    return traces


def np_2_tensor(phaseRad):
    _tensor = torch.from_numpy(phaseRad).float()
    _tensor = _tensor.unsqueeze(0).unsqueeze(0)  # 添加 batch 和 channel 维度
    return _tensor


def create_multibeam_dataset(bit_num, traces1, traces2):
    A1_tensor_list = []
    A2_tensor_list = []
    A_add_trad_tensor_list = []
    AM1_tensor_list = []
    AM2_tensor_list = []
    for trace1 in traces1:
        for trace2 in traces2:
            print(f"trace1={trace1}, trace2={trace2}")
            phaseRad_1 = point_2_phase(trace1[0], trace1[1])
            phaseRad_2 = point_2_phase(trace2[0], trace2[1])
            phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_complex_beam_n([phaseRad_1, phaseRad_2], bit_num)
            patternBit_dbw_1, patternBit_dbw_mask_1 = get_mask_by_slope(bit_num, phaseRad_1)
            patternBit_dbw_2, patternBit_dbw_mask_2 = get_mask_by_slope(bit_num, phaseRad_2)
            #
            A1_tensor_list.append(np_2_tensor(phaseRad_1))
            A2_tensor_list.append(np_2_tensor(phaseRad_2))
            AM1_tensor_list.append(np_2_tensor(patternBit_dbw_mask_1))
            AM2_tensor_list.append(np_2_tensor(patternBit_dbw_mask_2))
            #
            _tensor = torch.from_numpy(phase_mix).float()
            _tensor = _tensor.unsqueeze(0)  # 添加 batch 维度
            A_add_trad_tensor_list.append(_tensor)
            # draw_img(phases[0])
            # draw_img(phases[1])
            # draw_img(patternBit_dbw_mask_1)
            # draw_img(patternBit_dbw_mask_2)
            # print("-------------------------------------------------------")
    # 拼接为批量张量 (100, 1, 3, 3)
    A1_tensor_batch = torch.cat(A1_tensor_list, dim=0)
    A2_tensor_batch = torch.cat(A2_tensor_list, dim=0)
    A_add_trad_tensor_batch = torch.cat(A_add_trad_tensor_list, dim=0)
    AM1_tensor_batch = torch.cat(AM1_tensor_list, dim=0)
    AM2_tensor_batch = torch.cat(AM2_tensor_list, dim=0)
    # 输出形状验证
    print(f"[shape] A1:{A1_tensor_batch.shape}, A2:{A2_tensor_batch.shape}, "
          f"A_add_trad:{A_add_trad_tensor_batch.shape}, "
          f"AM1:{AM1_tensor_batch.shape}, AM2:{AM2_tensor_batch.shape}")  # 输出: torch.Size([100, 1, 3, 3])
    return A1_tensor_batch, A2_tensor_batch, A_add_trad_tensor_batch, AM1_tensor_batch, AM2_tensor_batch



# ----------------------------
# 1. 数据生成模块
# ----------------------------
def generate_data():
    bit_num = 1
    theta1_start, theta1_step, theta1_end, phi1_start, phi1_step, phi1_end = 30, 10, 35, 0, 20, 160
    theta2_start, theta2_step, theta2_end, phi2_start, phi2_step, phi2_end = 30, 10, 35, 180, 20, 340
    # 生成训练集数据
    traces1 = get_traces(theta1_start, theta1_step, theta1_end, phi1_start, phi1_step, phi1_end)
    traces2 = get_traces(theta2_start, theta2_step, theta2_end, phi2_start, phi2_step, phi2_end)
    # PS生成码阵数据集
    return create_multibeam_dataset(bit_num, traces1, traces2)






# ----------------------------
# 2. 模型定义模块
# ----------------------------
# 定义网络结构
# 定义网络结构
class MatrixAddNetwork(nn.Module):
    def __init__(self, input_shape):
        super(MatrixAddNetwork, self).__init__()
        self.input_shape = input_shape

        # 输入形状: (batch_size, 2, n, m)
        self.conv1 = nn.Conv2d(in_channels=2, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)

        # 计算展平后的大小
        conv1_output_shape = ((input_shape[0] // 2), (input_shape[1] // 2))
        conv2_output_shape = ((conv1_output_shape[0] // 2), (conv1_output_shape[1] // 2))
        flattened_size = 32 * conv2_output_shape[0] * conv2_output_shape[1]

        self.fc1 = nn.Linear(flattened_size, 128)
        self.fc2 = nn.Linear(128, input_shape[0] * input_shape[1])

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = self.pool(x)
        x = torch.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)  # 展平
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        # 将输出重塑为 (batch_size, n, m)
        return x.view(x.size(0), self.input_shape[0], self.input_shape[1])



# ----------------------------
# 3. 评估模块
# ----------------------------
# 定义损失函数
def loss_function(A_add):
    return torch.norm(A_add, p='fro')  # 可替换为你自己的解析函数


def train(model, inputs, targets, optimizer, criterion, num_epochs=1000, model_save_path="model.pt"):
    for epoch in range(num_epochs):
        model.train()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        loss_item = loss.item()
        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Train Loss: {loss_item:.6f}")
    # 训练结束后保存模型
    torch.save(model.state_dict(), model_save_path)
    print(f"Model saved to {model_save_path}")
    return model


def train_with_loss_function(model, inputs, AM1s, AM2s, optimizer, num_epochs=1000, model_save_path="model.pt"):
    for epoch in range(num_epochs):
        model.train()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function_mask(outputs, AM1s, AM2s)
        # loss.backward()
        optimizer.step()
        # loss_item = loss.item()
        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Train Loss: {loss:.6f}")
    # 训练结束后保存模型
    torch.save(model.state_dict(), model_save_path)
    print(f"Model saved to {model_save_path}")
    return model


def test(model, inputs, targets, criterion):
    model.eval()
    with torch.no_grad():
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    print(f"Test Loss: {loss.item():.6f}")
    return outputs, loss.item()

def test_with_loss_function(model, inputs, AM1s, AM2s):
    model.eval()
    with torch.no_grad():
        outputs = model(inputs)
        loss = loss_function_mask(outputs, AM1s, AM2s)
    print(f"Test Loss: {loss:.6f}")
    return outputs, loss


# ----------------------------
# 4. 主程序
# ----------------------------
def main(base_dir):
    input_shape = (64, 64)

    # 初始化网络和优化器
    net = MatrixAddNetwork(input_shape)
    optimizer = optim.Adam(net.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    # 生成数据
    A1_tensor, A2_tensor, A_add_true, AM1_tensor, AM2_tensor = generate_data()
    num_samples = A1_tensor.size(0)

    # 数据划分
    indices = torch.randperm(num_samples)
    train_ratio = 0.8
    train_size = int(train_ratio * num_samples)
    train_indices, test_indices = indices[:train_size], indices[train_size:]

    # 划分训练集与测试集
    A1_train, A1_test = A1_tensor[train_indices], A1_tensor[test_indices]
    A2_train, A2_test = A2_tensor[train_indices], A2_tensor[test_indices]
    A_add_train, A_add_test = A_add_true[train_indices], A_add_true[test_indices]
    AM1_train, AM1_test = AM1_tensor[train_indices], AM1_tensor[test_indices]
    AM2_train, AM2_test = AM2_tensor[train_indices], AM2_tensor[test_indices]

    # 构建输入张量
    inputs_train = torch.cat([A1_train, A2_train], dim=1)
    inputs_test = torch.cat([A1_test, A2_test], dim=1)

    # 训练
    train(net, inputs_train, A_add_train, optimizer, criterion, num_epochs=1000, model_save_path=base_dir+"model.pt")
    # train_with_loss_function(net, inputs_train, AM1_train, AM2_train, optimizer,
    #                          num_epochs=1000, model_save_path=base_dir+"model.pt")

    # 测试
    outputs, test_loss = test(net, inputs_test, A_add_test, criterion)
    # outputs, test_loss = test_with_loss_function(net, inputs_test, AM1_test, AM2_test)
    print(f"Final Test Loss: {test_loss:.6f}")

    # 展示测试结果
    # 将张量移动到 CPU 并转换为 numpy 数组
    outputs_np = outputs.numpy()
    A_add_test_np = A_add_test.numpy()

    # 创建一个包含 8 个 numpy 数组的列表
    outputs_np_list = [outputs_np[i] for i in range(outputs.size(0))]
    A_add_test_np_list = [A_add_test_np[i] for i in range(A_add_test.size(0))]

    # 打印结果以验证
    for i in range(len(outputs_np_list)):
        output = outputs_np_list[i]
        A_add_test = A_add_test_np_list[i]
        print(f"Array {i} shape: output={output.shape}, A_add_test={A_add_test.shape}")
        pattern_output = get_patterndbw(output)
        pattern_A_add_test = get_patterndbw(A_add_test)
        draw_img(A_add_test)
        draw_img(pattern_A_add_test)
        draw_img(output)
        draw_img(pattern_output)




if __name__ == "__main__":
    # 初始化日志
    # logger = setup_logging(base_path + "/trace.txt")
    logger = setup_logging()
    # 示例日志记录
    base_dir = "../files/dissertation/chapter_4/test_self_supervised/"
    logger.info(f"base_dir:{base_dir}")

    main(base_dir)